from pytorch_lightning import Trainer
from argparse import ArgumentParser
import torch
from tasks.ClassSummary.models.CsGen.model import CsGen
from tasks.ClassSummary.models.CsGen.data import CsGenDataModule
from config.Config import BASE_DIR
from pytorch_lightning import seed_everything


def add_program_args(parser: ArgumentParser):
    parser.add_argument("--code_field", help="代码词表文件", type=str,
                        default=str(BASE_DIR / "tasks/ClassSummary/resource/code-s.field"))
    parser.add_argument("--nl_field", help="自然语言词表文件", type=str,
                        default=str(BASE_DIR / "tasks/ClassSummary/resource/nl-s.field"))
    parser.add_argument("--vocab_map", help="词表映射文件", type=str,
                        default=str(BASE_DIR / "tasks/ClassSummary/resource/vocab_map.mat"))
    parser.add_argument("--check_point", help="check point path",
                        default=None)
    parser.add_argument("--save_path", help="结果存储文件", default=None)
    parser.add_argument("--seed", help="random seed", default=1024)
    parser.add_argument("--gpu_num", help="GPUS", default=0)


def main(args):
    seed_everything(args.seed)
    cf = torch.load(args.code_field)
    nf = torch.load(args.nl_field)
    vocab_map = torch.load(args.vocab_map)
    print("Loading vocab_map")
    vocab_map = vocab_map.to_dense().to(torch.float)
    if torch.cuda.is_available():
        vocab_map = vocab_map.to('cuda')
    print(vocab_map.device)
    print("Loaded vocab_map")
    data = CsGenDataModule(
        train_path=None,
        val_path=None,
        test_path=args.test,
        code_field=cf,
        nl_field=nf,
        batch_size=args.batch_size,
    )
    model = CsGen.load_from_checkpoint(args.check_point,
                                       vocab_map=vocab_map,
                                       code_vocab=cf.vocab,
                                       nl_vocab=nf.vocab,
                                       translate_path=args.save_path)
    trainer = Trainer(enable_checkpointing=False, gpus=args.gpu_num)

    trainer.test(model, datamodule=data)


if __name__ == '__main__':
    # 解析命令行参数
    arg_parser = ArgumentParser()
    add_program_args(arg_parser)
    CsGenDataModule.add_data_args(arg_parser)
    cli_args = arg_parser.parse_args()
    main(cli_args)
